import numpy as np
import os
import output_function as of

# 使用示例
base_path = './data/RB_perturb/'  # 根据您的保存路径进行修改

# 首先，只加载一次初始数据和模拟参数
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra_star = config['Ra_star']
Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 假设我们从out_iter=0开始，直到8000，且每次增加100
for out_iter in range(0, 6001, 100):
    # 加载数据
    n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
    
    # 确保加载了正确的时间步长
    if out_iter_loaded != out_iter:
        print(f"Failed to load the correct time step {out_iter} or file does not exist.")
        continue  # 如果加载失败，跳过当前迭代，继续下一个时间步
    
    # 打印一些变量以确认它们已被正确加载和赋值
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    
    # 使用加载的数据进行分析、绘图或其他处理
    base_folder=f'./RB_perturb_plot_Ra_{Ra_star}_Lx_{Lx}_miniter'
    of.diag_n_pcolor_filepath(n, n_init0_matrix, u, v, xx, yy, out_iter, Ra_star, dt,base_folder)
    of.plot_and_save_n_density_x_section(x,y, n, n_init0_matrix, Ny//2, out_iter,title_base="delta n_density Distribution",save_folder_base=base_folder)
    of.plot_and_save_v_x_section(x,y, v, Ny//2, out_iter,title_base="V Distribution along X-axis",save_folder_base=base_folder)